#include "cotask.h"
#include "context.h"

#include <stdlib.h>
#include <string.h>

#define DEFAULT_TASK_SIZE (1024 * 1024 * 2)
#define TASK_GUARD_SIZE (4 * sizeof(uint32_t))
#define TASK_MAGIC_HEAD1 (0x99008800)
#define TASK_MAGIC_HEAD2 (0x19490917)
#define TASK_MAGIC_HEAD3 (0xdeadbeef)
#define TASK_MAGIC_HEAD4 (0xabadabed)

typedef struct CoEnv
{
    CoTask *mainTask;
} CoEnv;

static __thread CoEnv gCoEnv = {NULL};

static const CoTaskAttr gDefaultAttr = {.stackSize = DEFAULT_TASK_SIZE};

void *CoSwapContext(Context *src, Context *dst, void *para) asm("Switch");
void CoSetTaskStatus(CoTask *task, CoTaskStatus status);
void CoTaskEntryFunction(CoTask *task, void *arg);

typedef struct CoTask
{
    Context context;
    CoTask *parent;
    ProcessFunction process;
    void *arg;
    CoTaskStatus status;
    uintptr_t stack;
} CoTask;

bool CoInitEnvironment()
{
    if (gCoEnv.mainTask != NULL)
    {
        return true;
    }

    CoTask *task = (CoTask *)malloc(sizeof(CoTask));
    if (task == NULL)
    {
        return false;
    }

    memset(task, 0, sizeof(CoTask));
    gCoEnv.mainTask = task;
    return true;
}

CoTask *CoCreateTask(const CoTaskAttr *attr, ProcessFunction process, void *arg)
{
    if (process == NULL)
    {
        return NULL;
    }

    CoInitEnvironment();
    attr = (attr == NULL) ? &gDefaultAttr : attr;
    uint32_t stackSize = (attr->stackSize < DEFAULT_TASK_SIZE) ? DEFAULT_TASK_SIZE : attr->stackSize;

    int allocSize = sizeof(CoTask) + stackSize;
    CoTask *task = (CoTask *)malloc(allocSize + TASK_GUARD_SIZE);
    if (task == NULL)
    {
        return NULL;
    }

    memset(task, 0, allocSize + TASK_GUARD_SIZE);
    uintptr_t stack = (uintptr_t)(allocSize - (sizeof(void *) << 1) + (uintptr_t)task);
    stack = (stack >> 4) << 4;

    ((uint32_t *)stack)[0] = TASK_MAGIC_HEAD1;
    ((uint32_t *)stack)[1] = TASK_MAGIC_HEAD2;
    ((uint32_t *)stack)[2] = TASK_MAGIC_HEAD3;
    ((uint32_t *)stack)[3] = TASK_MAGIC_HEAD4;

    task->process = process;
    task->stack = stack;
    task->arg = arg;
    task->parent = gCoEnv.mainTask;
    task->status = CO_TASK_STATUS_INIT;
    return task;
}

void *CoResumeTask(CoTask *task, void *arg)
{
    if (task == NULL)
    {
        return NULL;
    }

    if (CoGetTaskStatus(task) == CO_TASK_STATUS_FINISH)
    {
        return NULL;
    }

    if (CoGetTaskStatus(task) == CO_TASK_STATUS_INIT)
    {
        MakeContext(&task->context, (void *)task->stack, CoTaskEntryFunction);
        arg = task;
    }

    CoSetTaskStatus(task, CO_TASK_STATUS_RUNNING);
    return CoSwapContext(&(task->parent->context), &task->context, arg);
}
void *CoExitTask(CoTask *task, void *arg)
{
    CoSetTaskStatus(task, CO_TASK_STATUS_FINISH);
    return CoSwapContext(&task->context, &task->parent->context, arg);
}

void *CoYieldTask(CoTask *task, void *arg)
{
    CoSetTaskStatus(task, CO_TASK_STATUS_SUSPENDING);
    return CoSwapContext(&task->context, &task->parent->context, arg);
}
void CoDestroyTask(CoTask *task)
{
    if (task == NULL)
    {
        return;
    }
    free(task);
}

CoTaskStatus CoGetTaskStatus(CoTask *task)
{
    return (task != NULL) ? task->status : CO_TASK_STATUS_INVALID;
}

void CoSetTaskStatus(CoTask *task, CoTaskStatus status)
{
    if (task == NULL)
    {
        return;
    }
    task->status = status;
}

void CoTaskEntryFunction(CoTask *task, void *arg)
{
    void *result = task->process(task, task->arg);
    CoExitTask(task, result);
    return;
}